from pyabb import XYZPlatform
from pyati import ATISensor
from gelslimpi import GelSlimInterface, InputSource
from gelslimpi.visualize import qtcv

import cv2
import sys
from PyQt5 import QtWidgets
from scipy.interpolate import LinearNDInterpolator
from scipy.interpolate import RegularGridInterpolator as RGI

from datetime import datetime
from log_writer import LogWriter

from time import sleep
from threading import Thread
import numpy as np
from pathlib import Path
PROJ_DIR = Path(__file__).resolve().parent.parent
CONF_DIR = PROJ_DIR / "configs"



class DataCollector():
    """移动到任意像素坐标
    采集图片
    采集传感器数据
    下压"""
    def __init__(self, log_dir=PROJ_DIR / f"runs/{datetime.now().strftime('%y_%m_%d_%H_%M_%S')}"):
        self.xyz = XYZPlatform("/dev/ttyUSB0", [-57, -119, -48])
        self.ati = ATISensor()
        self.ati.tare()
        self.source = InputSource("/raspicam_node32/image/compressed", "gelslim")
        self.gelslim = GelSlimInterface(configs=CONF_DIR/"configs.yaml", input_source=self.source)
        self.init_interp()


        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(exist_ok=True, parents=True)
        self.log_writer = LogWriter(self.log_dir/"log.csv",
                                    fields=["cls_name", "id", "x_pix", "y_pix", "depth_mm",
                                            "fx", "fy", "fz", "tx", "ty", "tz", "x_mm", "y_mm", "z_mm"])
        self.z_cont = -38
        self.cont_th1 = -0.015
        self.cont_th2 = -0.035

    def show_img(self):
        while True:
            img = self.get_frame()
            qtcv.imshow("gelslim", img)
            key = qtcv.waitKey(200)
            if key == qtcv.Key.Key_Escape:
                break

    def get_frame(self):
        return self.gelslim.get_frame()[0]

    def get_ati(self):
        return self.ati.data

    def get_xyz(self):
        return self.xyz.x, self.xyz.y, self.xyz.z

    def init_interp(self, points_pix=None, points_xyz=None):
        """将像素坐标映射为实际坐标"""
        if points_pix is None:
            points_pix = np.array([[53, 50], [321, 50], [590, 50], [590, 240], [594, 432], [321, 429], [61, 430], [51, 239], [321, 239], [559, 399]])
        else:
            points_pix = np.array(points_pix)
        if points_xyz is None:
            points_xyz = np.array([[23.191, -34.53], [14.598, -34.585], [6.347, -34.506], [6.188, -28.976], [6.386, -23.036], [14.598, -22.989],
                                   [22.591, -23.055], [23.038, -28.797], [14.615, -28.756], [7.265, -23.931]])
        else:
            points_xyz = np.array(points_xyz)
        # 用非结构化插值获得网格数据
        interp = LinearNDInterpolator(points=points_pix, values=points_xyz)
        ymin = np.max(points_pix[0:3,1])
        xmax = np.min(points_pix[2:5,0])
        ymax = np.min(points_pix[4:7,1])
        xmin = np.max([points_pix[6,0], points_pix[7,0], points_pix[0,0]])
        x_span = np.linspace(xmin, xmax, 9)
        y_span = np.linspace(ymin, ymax, 9)
        X, Y = np.meshgrid(x_span, y_span, indexing="ij")
        grids = interp(X, Y)
        self.interp = RGI((x_span, y_span), grids, bounds_error=False, fill_value=None)

    def move(self, pix_x, pix_y):
        # 移动到接触高度上方， 力传感器清零
        self.xyz.move_abs(z = self.z_cont+2, vel=30)  # z 提升到 -26
        self.ati.tare()
        # xy 移动
        x, y = self.interp((pix_x, pix_y))
        print(f"move to [{pix_x:.2f}, {pix_y:.2f}]pix, [{x:.2f}, {y:.2f}]mm")
        self.xyz.move_abs(x=x, y=y, vel=30)
        # 向下移动
        self.xyz.move_rel(z=-1.3)  # for speed
        while not self.is_contacted(threshold=self.cont_th1):
            self.xyz.move_rel(z=-0.01, vel=20)
            sleep(0.05)
            img = self.get_frame()
            qtcv.imshow("contact", img)
            qtcv.waitKey(5)
        # 确认接触
        self.xyz.move_rel(z=-0.2, vel=30)
        if self.get_ati()[2] > self.cont_th2:
            return False
        else:
            self.z_cont = self.xyz.z + 0.2  # 更新上次接触的高度
            print(f"z_cont: {self.z_cont}")
            return True

    def log_data(self, cls_name, id, px, py, z_mm):
        """记录数据"""
        img = self.get_frame()
        img_path = str(self.log_dir / f"{cls_name}_{id:>04d}.jpg")
        cv2.imwrite(img_path, img)
        data = self.ati.data
        self.log_writer.log(cls_name, id, px, py, z_mm, *list(data),
                            self.xyz.x, self.xyz.y, self.xyz.z, echo=True)
        qtcv.imshow("contact", img)
        qtcv.waitKey(5)

    def is_contacted(self, threshold):
        """检测接触"""
        fz = self.get_ati()[2]
        # print(fz)
        return fz <= threshold


if __name__ == '__main__':
    cnt = 0
    id = 0
    cls_name = "sphere6"
    dc = DataCollector()
    while cnt < 200:
        x = np.random.randint(20, 620)
        y = np.random.randint(20, 460)
        depths = np.random.uniform(0.2, 1.1, 2)
        ret = dc.move(x, y)
        if not ret:
            continue
        cnt += 1
        z = dc.z_cont
        for depth in depths:
            dc.xyz.move_abs(z=z-depth)
            sleep(0.5)
            dc.log_data(cls_name, id, x, y, depth)

            id += 1
    dc.xyz.move_abs(z=-22)

